import cv2
import numpy
from insightface.app import FaceAnalysis
from insightface.model_zoo.inswapper import INSwapper
# from torchvision import transforms
from typing import Any

from Settings import ModelsBaseDIR, OutBaseDir
from core.face.Face import Face


class FaceEngine:
    def __init__(self):
        # 模型地址

        # 人脸识别
        self.__face_analysis = FaceAnalysis(
            # 下载模型的目录名称
            name='buffalo_l',
            # 模型下载位置
            root=ModelsBaseDIR,
            # root=self.__face_model_path,
            # 'AzureExecutionProvider, CPUExecutionProvider'
            providers=["CPUExecutionProvider"]
        )
        # 预处理
        self.__face_analysis.prepare(ctx_id=0, det_size=(640, 640))

        self.initInSwapper()

    def initInSwapper(self):
        self.__inswapper = INSwapper(
            model_file=fr'C:\Users\Administrator\Desktop\change_face\change_face\models\1.onnx',
        )

    def savePicture(self, enhance_face_image):
        pass
        # pil_image = transforms.ToPILImage()(enhance_face_image[0])
        # result = pil_image.save(f'{OutBaseDir}/output_image.jpg')
        # print(f"保存完毕，输出{result}")

    def get_one_face(self, img_path):
        frame = cv2.imread(img_path)
        face = self.__face_analysis.get(frame)
        try:
            # 如果拿到了多张脸，就只返回最近的一个
            return min(face, key=lambda x: x.bbox[0])
        except ValueError:
            raise "传的数据格式不对"

    def swap_face(self,tmp_face: numpy.ndarray[Any, Any], target_face: Face, source_face: Face):
        # self.__face_model_path = f'{ModelsBaseDIR}/inswapper_128_fp16.onnx'
        # self.__face_model_path = fr'C:\Users\Administrator\Desktop\change_face\change_face\models\inswapper_128_fp16.onnx'
        # # 加载训练好的模型
        # face_swapper:INSwapper = insightface.model_zoo.get_model(
        #     name=self.__face_model_path,
        #     # 'AzureExecutionProvider, CPUExecutionProvider'
        #     providers=["CPUExecutionProvider"]
        # )
        # frame = face_swapper.get(tmp_face, target_face, source_face, paste_back=True)
        frame = self.__inswapper.get(tmp_face, target_face, source_face, paste_back=True)
        return frame

    def _pointFace(self, source_image_path, region_result: Face):
        landmark = region_result['landmark_2d_106']
        print(landmark)
        # print(region_result.keys())

        color = (0, 0, 255)
        tim = cv2.imread(source_image_path)
        for face_landmark in landmark:
            # 不要问为什么写的这么屎，仅仅就是给下面标点参数格式化下
            circle_index = tuple(int(round(i)) for i in face_landmark)

            cv2.circle(tim, circle_index, 1, color, 1, cv2.LINE_AA)
        cv2.imwrite(f'{OutBaseDir}/test_out1.jpg', tim)


if __name__ == '__main__':
    faceEngine = FaceEngine()
    from Settings import ModelsBaseDIR, OutBaseDir,BaseTestDataCubePath
    source_image_path = f'{BaseTestDataCubePath}/4.jpg'
    source_region_result: Face = faceEngine.get_one_face(source_image_path)
    #
    target_image_path = f'{BaseTestDataCubePath}/1.jpg'
    target_region_result: Face = faceEngine.get_one_face(target_image_path)

    tmp_region_result: numpy.ndarray[Any, Any] = cv2.imread(target_image_path)

    tmp_result = faceEngine.swap_face(tmp_region_result, target_region_result, source_region_result)
    print(tmp_result)

    cv2.imwrite(f'{OutBaseDir}/test_out3.jpg', tmp_result)